import torch
import torchvision

def test_cuda_and_torch():
    # 检查 CUDA 是否可用
    if torch.cuda.is_available():
        print("CUDA is available.")
        print(f"CUDA version: {torch.version.cuda}")
        print(f"cuDNN version: {torch.backends.cudnn.version()}")

        # 获取 GPU 信息
        device = torch.cuda.current_device()
        gpu_name = torch.cuda.get_device_name(device)
        gpu_memory = torch.cuda.get_device_properties(device).total_memory / (1024 ** 3)  # 转换为 GB

        print(f"GPU Name: {gpu_name}")
        print(f"GPU Memory: {gpu_memory:.2f} GB")

        # 测试 PyTorch 和 TorchVision 是否可用
        print(f"PyTorch version: {torch.__version__}")
        print(f"TorchVision version: {torchvision.__version__}")

        # 创建一个张量并将其移动到 GPU 上
        tensor = torch.randn(3, 3).cuda()
        print("Tensor on GPU:")
        print(tensor)
    else:
        print("CUDA is not available.")

if __name__ == "__main__":
    test_cuda_and_torch()
